package com.gitee.gateway.core.security;

import com.gitee.gateway.core.cache.CacheMgr;
import io.github.resilience4j.core.ConfigurationNotFoundException;
import io.github.resilience4j.ratelimiter.RateLimiter;
import io.github.resilience4j.ratelimiter.RateLimiterConfig;
import io.github.resilience4j.ratelimiter.internal.AtomicRateLimiter;
import io.github.resilience4j.ratelimiter.internal.InMemoryRateLimiterRegistry;

import java.util.Map;
import java.util.Objects;
import java.util.function.Supplier;

/**
 * 使用缓存来管理 RateLimiter 实例，解决的问题：
 * 1. 分布式环境下的使用（尚无法共享状态）
 * 2. 避免内存过载
 */
public class CacheRateLimiterRegistry extends InMemoryRateLimiterRegistry {

    private final static String REGION = "rateLimiter";
    private final String catalog;

    public CacheRateLimiterRegistry(String catalog, RateLimiterConfig defaultConfig) {
        super(defaultConfig);
        this.catalog = catalog;
    }

    @Override
    public RateLimiter rateLimiter(String name, RateLimiterConfig rateLimiterConfig, Map<String, String> tags) {
        String key = catalog + name;
        return CacheMgr.get(REGION, key, () -> new AtomicRateLimiter(name,
                Objects.requireNonNull(rateLimiterConfig, CONFIG_MUST_NOT_BE_NULL), getAllTags(tags)));
    }

    @Override
    public RateLimiter rateLimiter(String name, Supplier<RateLimiterConfig> rateLimiterConfigSupplier, Map<String, String> tags) {
        String key = catalog + name;
        return CacheMgr.get(REGION, key, () -> new AtomicRateLimiter(name, Objects.requireNonNull(
                Objects.requireNonNull(rateLimiterConfigSupplier, SUPPLIER_MUST_NOT_BE_NULL).get(),
                CONFIG_MUST_NOT_BE_NULL), getAllTags(tags)));
    }

    @Override
    public RateLimiter rateLimiter(String name, String configName, Map<String, String> tags) {
        String key = catalog + name;
        return CacheMgr.get(REGION, key, () -> RateLimiter.of(name, getConfiguration(configName)
                .orElseThrow(() -> new ConfigurationNotFoundException(configName)), getAllTags(tags)));
    }
}
